自定义层¶

Note

有时候我们想要一个层实现某种功能，但tensorflow没有其默认实现，在这种情况下就需要自定义层。

无参数层¶

import tensorflow as tf
from tensorflow import keras

# 使用keras.layers.Lambda
exponential_layer = keras.layers.Lambda(lambda x: tf.exp(x))
# as expected
exponential_layer([0., 1.])

<tf.Tensor: shape=(2,), dtype=float32, numpy=array([1.       , 2.7182817], dtype=float32)>

keras.layers.Layer的子类¶

最好的自定义层的方式是创建一个 keras.layers.Layer 的子类。

class MyDense(tf.keras.layers.Layer):
    def __init__(self, num_outputs):
        # 与input_shape无关的initialization
        super(MyDense, self).__init__()
        self.num_outputs = num_outputs
        
    def build(self, input_shape):
        # 与input_shape相关的initialization
        # 首次使用层时会被调用
        self.kernel = self.add_weight("kernel",
                                      shape=[int(input_shape[-1]),
                                             self.num_outputs])

    def call(self, inputs):
        # 正向传播
        return tf.matmul(inputs, self.kernel)

# 像其他任何层一样使用MyDense层
model = keras.models.Sequential([
    MyDense(30),
    MyDense(1)
])

自定义自定义模型